"""Command line utilities for external Codex/Claude clients.

This CLI bridges Codex/Claude terminals into the Agent-Society presence API.
It provides simple login/heartbeat/logout commands plus higher-level helpers
that wrap `codex -p` executions, capture a timeline, and optionally orchestrate
multiple roles under a coordinating "system" agent.
"""
from __future__ import annotations

import argparse
import datetime as dt
import json
import os
import shutil
import subprocess
import sys
import textwrap
import threading
import time
from dataclasses import dataclass, asdict
from pathlib import Path
from typing import Any, Callable, Iterable

import httpx
from agent_society.memory.event_log import append_event

DEFAULT_BASE_URL = os.environ.get("AGENT_SOCIETY_API", "http://localhost:8080")
CONFIG_HOME = Path(os.environ.get("AGENT_SOCIETY_HOME", Path.home() / ".agent-society"))
SESSION_PATH = CONFIG_HOME / "session.json"
DEFAULT_TIMELINE = Path(os.environ.get("AGENT_SOCIETY_TIMELINE", "data/codex_timeline.jsonl"))
AGENT_CONFIG_PATH = Path(os.environ.get("AGENT_CONFIG_PATH", "config/agents.json"))
AGENT_EVENT_DIR = Path(os.environ.get("AGENT_EVENT_DIR", "events"))
PERSONA_SUMMARY_PATH = Path(os.environ.get("AGENT_PERSONA_SUMMARIES", "data/persona_summaries.json"))


class CLIError(RuntimeError):
    """Raised when a command fails for a user-facing reason."""


@dataclass
class SessionState:
    base_url: str
    agent_id: str
    client: str
    session_id: str
    capabilities: list[str]
    token: str | None = None

    def to_dict(self) -> dict[str, Any]:
        return {
            "base_url": self.base_url,
            "agent_id": self.agent_id,
            "client": self.client,
            "session_id": self.session_id,
            "capabilities": self.capabilities,
            "token": self.token,
        }

    @classmethod
    def load(cls) -> "SessionState | None":
        if not SESSION_PATH.exists():
            return None
        data = json.loads(SESSION_PATH.read_text(encoding="utf-8"))
        return cls(
            base_url=data["base_url"],
            agent_id=data["agent_id"],
            client=data.get("client", "external"),
            session_id=data["session_id"],
            capabilities=data.get("capabilities", []),
            token=data.get("token"),
        )

    def save(self) -> None:
        CONFIG_HOME.mkdir(parents=True, exist_ok=True)
        SESSION_PATH.write_text(json.dumps(self.to_dict(), indent=2, ensure_ascii=False), encoding="utf-8")

    @classmethod
    def delete(cls) -> None:
        if SESSION_PATH.exists():
            SESSION_PATH.unlink()


# ---------------------------------------------------------------------------
# API helpers
# ---------------------------------------------------------------------------


def client_for(base_url: str, token: str | None = None) -> httpx.Client:
    headers = {"Content-Type": "application/json"}
    if token:
        headers["Authorization"] = f"Bearer {token}"
    return httpx.Client(base_url=base_url.rstrip("/"), headers=headers, timeout=10.0)


def _resolve_token(arg_token: str | None) -> str | None:
    return arg_token or os.environ.get("AGENT_SOCIETY_TOKEN")


def login_session(
    *,
    base_url: str,
    agent: str,
    client_name: str,
    capabilities: Iterable[str] | None,
    token: str | None,
) -> SessionState:
    caps = list(capabilities or [])
    with client_for(base_url, token) as client:
        response = client.post(
            "/agents/session",
            json={
                "agentId": agent,
                "client": client_name,
                "capabilities": caps,
            },
        )
        if response.status_code == 404:
            raise CLIError(
                f"Agent '{agent}' is not registered. Update config/agents.json or set AGENT_CONFIG_PATH."
            )
        if response.status_code == 409:
            raise CLIError("Agent already online. Run `agent-society logout` first or close stale sessions.")
        response.raise_for_status()
        data = response.json()
    return SessionState(
        base_url=base_url,
        agent_id=data["agentId"],
        client=client_name,
        session_id=data["sessionId"],
        capabilities=caps,
        token=token,
    )


def heartbeat_session(
    session: SessionState,
    *,
    state: str,
    capabilities: Iterable[str] | None = None,
) -> dict[str, Any]:
    payload = {"state": state}
    if capabilities is not None:
        payload["capabilities"] = list(capabilities)
    with client_for(session.base_url, session.token) as client:
        response = client.post(f"/agents/session/{session.session_id}/heartbeat", json=payload)
        response.raise_for_status()
        data = response.json()
    if capabilities is not None:
        session.capabilities = list(capabilities)
    return data


def logout_session(session: SessionState) -> dict[str, Any] | None:
    with client_for(session.base_url, session.token) as client:
        response = client.delete(f"/agents/session/{session.session_id}")
        if response.status_code == 404:
            return None
        response.raise_for_status()
        return response.json()


def list_presence(base_url: str, token: str | None) -> list[dict[str, Any]]:
    with client_for(base_url, token) as client:
        response = client.get("/agents/presence")
        response.raise_for_status()
        return response.json()


def load_agent_registry() -> list[dict[str, Any]]:
    if not AGENT_CONFIG_PATH.exists():
        raise CLIError(
            f"Agent registry not found at {AGENT_CONFIG_PATH}. Set AGENT_CONFIG_PATH or create the file."
        )
    try:
        data = json.loads(AGENT_CONFIG_PATH.read_text(encoding="utf-8"))
    except json.JSONDecodeError as exc:
        raise CLIError(f"Failed to parse agent registry: {exc}") from exc
    if not isinstance(data, list):
        raise CLIError("Agent registry must be a JSON list of agent descriptors.")
    return data


def load_agent_profile(agent_id: str) -> dict[str, Any]:
    for entry in load_agent_registry():
        if entry.get("id") == agent_id:
            return entry
    raise CLIError(f"Agent '{agent_id}' not found in {AGENT_CONFIG_PATH}.")


def load_persona_summaries() -> dict[str, Any]:
    if not PERSONA_SUMMARY_PATH.exists():
        return {}
    try:
        data = json.loads(PERSONA_SUMMARY_PATH.read_text(encoding="utf-8"))
    except json.JSONDecodeError:
        return {}
    if not isinstance(data, dict):
        return {}
    return data


def save_persona_summary(agent_id: str, payload: dict[str, Any]) -> None:
    payload = dict(payload)
    existing = load_persona_summaries()
    existing[agent_id] = payload
    PERSONA_SUMMARY_PATH.parent.mkdir(parents=True, exist_ok=True)
    PERSONA_SUMMARY_PATH.write_text(json.dumps(existing, indent=2, ensure_ascii=False), encoding="utf-8")


def fetch_memories(
    *,
    base_url: str,
    agent: str,
    token: str | None,
    limit: int,
    query: str | None,
) -> list[dict[str, Any]]:
    params = {"limit": str(limit)}
    if query:
        params["query"] = query
    with client_for(base_url, token) as client:
        response = client.get(f"/agents/{agent}/memories", params=params)
        response.raise_for_status()
        return response.json()


def fetch_environment_view(base_url: str, token: str | None, limit: int) -> dict[str, Any]:
    params = {"limit": str(limit)}
    with client_for(base_url, token) as client:
        response = client.get("/environment/tree", params=params)
        response.raise_for_status()
        return response.json()


def format_memories_for_prompt(memories: list[dict[str, Any]]) -> str:
    lines: list[str] = []
    for idx, record in enumerate(memories, start=1):
        created = record.get("createdAt", "")
        content = (record.get("content") or "").replace("\n", " ")
        score = record.get("score", 0.0)
        lines.append(f"{idx}. [{created}] score={score:.2f} :: {content}")
    return "\n".join(lines) if lines else "(no recent memories)"


def build_reflection_prompt(
    *,
    agent_profile: dict[str, Any],
    memories: list[dict[str, Any]],
    previous_summary: dict[str, Any] | None,
) -> str:
    name = agent_profile.get("name", agent_profile.get("id"))
    role = agent_profile.get("role", "unspecified")
    objectives = agent_profile.get("objectives", [])
    traits = agent_profile.get("traits", {})
    objectives_line = ", ".join(objectives) if objectives else "(none listed)"
    traits_line = ", ".join(f"{k}: {v}" for k, v in traits.items()) if traits else "(none listed)"
    previous = previous_summary.get("summary") if previous_summary else None
    previous_line = previous if previous else "(no previous persona summary)"
    memory_block = format_memories_for_prompt(memories)
    prompt = f"""
You are the reflection subsystem for Agent-Society. Update the persona digest for the agent below.

Agent ID: {agent_profile.get('id')}
Display Name: {name}
Role: {role}
Objectives: {objectives_line}
Traits: {traits_line}
Last Known Summary: {previous_line}

Recent high-signal memories:
{memory_block}

Respond with a compact JSON object using these keys:
  "summary" (<=180 words) – updated narrative persona.
  "core_characteristics" (<=120 words) – identity, values, capabilities.
  "current_focus" (<=120 words) – near-term goals and concerns.
  "recent_insights" (<=100 words) – new learnings from the memories above.
  "next_steps" (<=80 words) – recommended follow-up actions or watchpoints.

Keep values plain text (no markdown, quotes, or bullet prefixes). Avoid newline characters inside values.
"""
    return textwrap.dedent(prompt).strip()


# ---------------------------------------------------------------------------
# Timeline helpers
# ---------------------------------------------------------------------------


def append_timeline_entry(path: Path, entry: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("a", encoding="utf-8") as fh:
        fh.write(json.dumps(entry, ensure_ascii=False) + "\n")


def summarise_timeline(entries: list[dict[str, Any]], limit: int = 8) -> str:
    recent = entries[-limit:]
    lines: list[str] = []
    for item in recent:
        ts = item.get("timestamp")
        agent = item.get("agent")
        summary = (item.get("output") or "").strip().splitlines()
        snippet = summary[0][:120] if summary else "(no output)"
        lines.append(f"[{ts}] {agent}: {snippet}")
    return "\n".join(lines) if lines else "(no prior activity)"


def fetch_timeline_entries(base_url: str, token: str | None, limit: int) -> list[dict[str, Any]]:
    params = {"limit": str(limit)}
    with client_for(base_url, token) as client:
        response = client.get("/timeline", params=params)
        response.raise_for_status()
        return response.json()


def format_timeline_entry(entry: dict[str, Any]) -> str:
    timestamp = entry.get("timestamp", "?")
    agent = entry.get("agent", "?")
    role = entry.get("role")
    role_hint = f" ({role})" if role and role != agent else ""
    entry_type = entry.get("type", "agent")
    duration = entry.get("duration")
    duration_hint = f" {duration:.1f}s" if isinstance(duration, (int, float)) else ""
    snippet_source = entry.get("output") or entry.get("prompt") or ""
    snippet = snippet_source.strip().splitlines()[0] if snippet_source else ""
    if len(snippet) > 90:
        snippet = snippet[:87] + "…"
    return f"[{timestamp}] {agent}{role_hint} [{entry_type}]{duration_hint}: {snippet}"


# ---------------------------------------------------------------------------
# Codex execution helpers
# ---------------------------------------------------------------------------


def ensure_codex_binary(binary: str) -> str:
    resolved = shutil.which(binary)
    if not resolved:
        raise CLIError(f"Unable to find codex binary '{binary}'. Adjust --codex-bin or update PATH.")
    return resolved


def run_codex_prompt(binary: str, prompt: str, extra_args: list[str] | None = None) -> subprocess.CompletedProcess[str]:
    cmd = [binary, "-p", prompt]
    if extra_args:
        cmd.extend(extra_args)
    return subprocess.run(cmd, capture_output=True, text=True)


def execute_codex_interaction(
    *,
    agent: str,
    role: str,
    client_name: str,
    capabilities: list[str],
    prompt: str,
    base_url: str,
    token: str | None,
    codex_bin: str,
    state: str,
    heartbeat_interval: int,
    timeline_path: Path | None,
    entry_type: str,
    extra_args: list[str] | None = None,
) -> dict[str, Any]:
    session = login_session(
        base_url=base_url,
        agent=agent,
        client_name=client_name,
        capabilities=capabilities,
        token=token,
    )
    stop_event = threading.Event()
    heartbeat_thread: threading.Thread | None = None

    def heartbeat_loop() -> None:
        while not stop_event.is_set():
            try:
                heartbeat_session(session, state=state, capabilities=capabilities)
            except Exception as exc:  # pragma: no cover - best-effort logging only
                print(f"heartbeat error: {exc}", file=sys.stderr)
            finally:
                if stop_event.wait(heartbeat_interval):
                    break

    if heartbeat_interval > 0:
        heartbeat_thread = threading.Thread(target=heartbeat_loop, daemon=True)
        heartbeat_thread.start()

    codex_path = ensure_codex_binary(codex_bin)
    started_at = dt.datetime.utcnow()
    started_monotonic = time.monotonic()
    result: subprocess.CompletedProcess[str]
    try:
        result = run_codex_prompt(codex_path, prompt, extra_args)
    finally:
        stop_event.set()
        if heartbeat_thread:
            heartbeat_thread.join()
        logout_session(session)

    ended_at = dt.datetime.utcnow()
    duration = time.monotonic() - started_monotonic
    entry = {
        "timestamp": started_at.isoformat() + "Z",
        "ended_at": ended_at.isoformat() + "Z",
        "agent": agent,
        "role": role,
        "client": client_name,
        "type": entry_type,
        "prompt": prompt,
        "returncode": result.returncode,
        "duration": duration,
        "output": result.stdout,
        "stderr": result.stderr,
    }
    if timeline_path:
        append_timeline_entry(timeline_path, entry)
    return entry


# ---------------------------------------------------------------------------
# Command handlers
# ---------------------------------------------------------------------------


def _print_json(payload: Any) -> None:
    json.dump(payload, sys.stdout, indent=2, ensure_ascii=False)
    sys.stdout.write("\n")


def cmd_login(args: argparse.Namespace) -> None:
    token = _resolve_token(args.token)
    session = login_session(
        base_url=args.base_url,
        agent=args.agent,
        client_name=args.client,
        capabilities=args.capabilities,
        token=token,
    )
    session.save()
    print(f"Connected as {session.agent_id} ({session.client}); session {session.session_id}")


def cmd_heartbeat(args: argparse.Namespace) -> None:
    session = SessionState.load()
    if not session:
        raise CLIError("No active session. Run `agent-society login` first.")
    capabilities = args.capabilities or session.capabilities
    data = heartbeat_session(session, state=args.state, capabilities=capabilities)
    session.capabilities = list(capabilities)
    session.save()
    status = data["presence"]["state"]
    print(f"Heartbeat acknowledged ({status}); lastSeen={data['presence'].get('lastSeen')}")


def cmd_logout(args: argparse.Namespace) -> None:
    session = SessionState.load()
    if not session:
        print("Already logged out")
        return
    data = logout_session(session)
    if data:
        print(f"Disconnected: {data['agentId']} is now {data['presence']['state']}")
    else:
        print("Session already closed on server; removing local state")
    SessionState.delete()


def cmd_watch(args: argparse.Namespace) -> None:
    session = SessionState.load()
    if not session:
        raise CLIError("No active session. Run `agent-society login` first.")
    interval = args.interval
    capabilities = args.capabilities or session.capabilities
    print(f"Sending heartbeats every {interval}s. Press Ctrl+C to stop.")
    try:
        while True:
            data = heartbeat_session(session, state=args.state, capabilities=capabilities)
            session.capabilities = list(capabilities)
            session.save()
            print(
                f"[{time.strftime('%H:%M:%S')}] state={data['presence']['state']} "
                f"lastSeen={data['presence'].get('lastSeen')}"
            )
            time.sleep(interval)
    except KeyboardInterrupt:
        print("Stopped heartbeats")


def cmd_presence(args: argparse.Namespace) -> None:
    token = _resolve_token(args.token)
    records = list_presence(args.base_url, token)
    if args.format == "json":
        _print_json(records)
        return
    print(f"Presence @ {args.base_url}")
    print("AGENT       STATE      CLIENT        CAPABILITIES")
    for record in records:
        presence = record.get("presence", {})
        caps = ",".join(presence.get("capabilities") or []) or "-"
        client = presence.get("client") or "-"
        print(f"{record['agentId']:<11}{presence.get('state','offline'):<11}{client:<13}{caps}")


def cmd_timeline(args: argparse.Namespace) -> None:
    token = _resolve_token(args.token)
    base_url = args.base_url
    if args.follow:
        try:
            with client_for(base_url, token) as client:
                params = {"limit": str(args.limit)}
                with client.stream("GET", "/timeline/stream", params=params) as response:
                    response.raise_for_status()
                    for line in response.iter_lines():
                        if not line or line.startswith(":"):
                            continue
                        if not line.startswith("data:"):
                            continue
                        payload = line[5:].strip()
                        if not payload:
                            continue
                        try:
                            entry = json.loads(payload)
                        except json.JSONDecodeError:
                            continue
                        if args.format == "json":
                            _print_json(entry)
                        else:
                            print(format_timeline_entry(entry))
        except KeyboardInterrupt:
            print("Stopped timeline stream")
    else:
        entries = fetch_timeline_entries(base_url, token, args.limit)
        if args.format == "json":
            _print_json(entries)
            return
        for entry in entries:
            print(format_timeline_entry(entry))


def cmd_memories(args: argparse.Namespace) -> None:
    session = SessionState.load()
    agent = args.agent or (session.agent_id if session else None)
    if not agent:
        raise CLIError("Specify --agent or log in first.")
    base_url = args.base_url or (session.base_url if session else DEFAULT_BASE_URL)
    token = _resolve_token(args.token) or (session.token if session else None)
    records = fetch_memories(
        base_url=base_url,
        agent=agent,
        token=token,
        limit=args.limit,
        query=args.query,
    )
    if args.format == "json":
        _print_json(records)
        return
    print(f"Top memories for {agent} @ {base_url}")
    print("SCORE  IMP  REC  REL  CREATED               CONTENT")
    for record in records:
        created = record.get("createdAt", "")
        snippet = (record.get("content") or "").replace("\n", " ")[:80]
        print(
            f"{record.get('score', 0.0):>4.2f}  {record.get('importance', 0.0):>4.2f}"
            f"  {record.get('recency', 0.0):>4.2f}  {record.get('relevance', 0.0):>4.2f}"
            f"  {created:<20}  {snippet}"
        )


def cmd_reflect(args: argparse.Namespace) -> None:
    base_url = args.base_url or DEFAULT_BASE_URL
    token = _resolve_token(args.token)
    profile = load_agent_profile(args.agent)
    summaries = load_persona_summaries()
    previous = summaries.get(args.agent)
    memories = fetch_memories(
        base_url=base_url,
        agent=args.agent,
        token=token,
        limit=args.memories,
        query=args.query,
    )
    prompt = build_reflection_prompt(
        agent_profile=profile,
        memories=memories,
        previous_summary=previous,
    )
    capabilities = list(args.capabilities or ["reflection"])
    entry = execute_codex_interaction(
        agent=args.agent,
        role=args.role or profile.get("role", args.agent),
        client_name=args.client,
        capabilities=capabilities,
        prompt=prompt,
        base_url=base_url,
        token=token,
        codex_bin=args.codex_bin,
        state=args.state,
        heartbeat_interval=args.heartbeat_interval,
        timeline_path=Path(args.timeline) if args.timeline else None,
        entry_type="reflection",
        extra_args=args.extra,
    )
    if entry["returncode"] != 0:
        print(entry.get("stderr", entry["output"] or "codex failed"), file=sys.stderr)
        raise CLIError(f"codex exited with status {entry['returncode']}")
    try:
        parsed = extract_decision(entry["output"])
    except CLIError:
        parsed = {"summary": entry["output"].strip()}
    updated_at = dt.datetime.utcnow().replace(tzinfo=dt.timezone.utc).isoformat().replace("+00:00", "Z")
    summary_payload = {
        "agentId": args.agent,
        "updatedAt": updated_at,
        "summary": parsed.get("summary", entry["output"].strip()),
        "coreCharacteristics": parsed.get("core_characteristics"),
        "currentFocus": parsed.get("current_focus"),
        "recentInsights": parsed.get("recent_insights"),
        "nextSteps": parsed.get("next_steps"),
        "topMemories": [
            {
                "id": mem.get("id"),
                "score": mem.get("score"),
                "createdAt": mem.get("createdAt"),
                "content": mem.get("content"),
            }
            for mem in memories
        ],
        "prompt": prompt,
        "capabilities": capabilities,
        "source": "codex",
    }
    if not args.no_cache:
        save_persona_summary(args.agent, summary_payload)
    if not args.no_event:
        summary_lines = [
            f"Persona refresh @ {updated_at}",
            f"Summary: {summary_payload['summary']}",
        ]
        if summary_payload.get("currentFocus"):
            summary_lines.append(f"Current focus: {summary_payload['currentFocus']}")
        if summary_payload.get("nextSteps"):
            summary_lines.append(f"Next steps: {summary_payload['nextSteps']}")
        append_event(
            AGENT_EVENT_DIR,
            owner=args.agent,
            actor=args.agent,
            ev_type="reflection",
            content="\n".join(summary_lines),
            headers={
                "kind": "persona_summary",
                "updated_at": updated_at,
            },
        )
    if args.format == "json":
        _print_json(summary_payload)
    else:
        print(f"Persona summary for {args.agent} ({updated_at})\n")
        print(summary_payload["summary"])
        if summary_payload.get("currentFocus"):
            print(f"\nCurrent focus: {summary_payload['currentFocus']}")
        if summary_payload.get("nextSteps"):
            print(f"Next steps: {summary_payload['nextSteps']}")


def cmd_environment(args: argparse.Namespace) -> None:
    base_url = args.base_url or DEFAULT_BASE_URL
    token = _resolve_token(args.token)
    view = fetch_environment_view(base_url, token, args.limit)
    if args.format == "json":
        _print_json(view)
    else:
        print(view.get("context", "(no context returned)"))


def cmd_codex_run(args: argparse.Namespace) -> None:
    token = _resolve_token(args.token)
    entry = execute_codex_interaction(
        agent=args.agent,
        role=args.role or args.agent,
        client_name=args.client,
        capabilities=list(args.capabilities or []),
        prompt=args.prompt,
        base_url=args.base_url,
        token=token,
        codex_bin=args.codex_bin,
        state=args.state,
        heartbeat_interval=args.heartbeat_interval,
        timeline_path=Path(args.timeline) if args.timeline else None,
        entry_type="agent",
        extra_args=args.extra,
    )
    if entry["returncode"] != 0:
        print(entry["stderr"], file=sys.stderr)
        raise CLIError(f"codex exited with status {entry['returncode']}")
    print(entry["output"])


def load_schedule_config(path: Path) -> dict[str, Any]:
    text = path.read_text(encoding="utf-8")
    if path.suffix.lower() in {".yaml", ".yml"}:
        try:
            import yaml  # type: ignore
        except ImportError as exc:  # pragma: no cover - optional dependency
            raise CLIError("PyYAML not installed; install pyyaml to use YAML configs.") from exc
        return yaml.safe_load(text)
    return json.loads(text)


def extract_decision(raw_output: str) -> dict[str, Any]:
    raw_output = raw_output.strip()
    # Attempt to parse entire output as JSON first
    try:
        data = json.loads(raw_output)
        if isinstance(data, dict):
            return data
    except json.JSONDecodeError:
        pass
    # Fallback: locate first JSON object in text
    start = raw_output.find("{")
    end = raw_output.rfind("}")
    if start == -1 or end == -1 or end <= start:
        raise CLIError("System agent response did not contain a JSON object.")
    snippet = raw_output[start : end + 1]
    return json.loads(snippet)


def cmd_codex_orchestrate(args: argparse.Namespace) -> None:
    config = load_schedule_config(Path(args.config))
    steps = args.steps or config.get("iterations", 1)
    codex_bin = args.codex_bin or config.get("codex_bin", "codex")
    base_url = args.base_url or config.get("base_url", DEFAULT_BASE_URL)
    token = _resolve_token(args.token) or config.get("token")
    timeline_path = Path(args.timeline or config.get("timeline_path", DEFAULT_TIMELINE))
    heartbeat_interval = args.heartbeat_interval or config.get("heartbeat_interval", 90)
    context_limit = args.context or config.get("context_limit", 6)
    environment_limit = config.get("environment_limit", 5)

    system_cfg = config.get("system") or {}
    agents_cfg: dict[str, Any] = config.get("agents") or {}
    if not system_cfg:
        raise CLIError("Schedule config missing 'system' block.")
    if not agents_cfg:
        raise CLIError("Schedule config missing 'agents' block.")

    timeline: list[dict[str, Any]] = []

    for step in range(1, steps + 1):
        context = summarise_timeline(timeline, limit=context_limit)
        environment_view = fetch_environment_view(base_url, token, environment_limit)
        environment_context = environment_view.get("context", "")
        system_prompt_template = system_cfg.get("prompt")
        if not system_prompt_template:
            raise CLIError("System config requires a 'prompt' template.")
        system_prompt = system_prompt_template.format(
            context=context,
            step=step,
            environment=environment_context,
            environment_context=environment_context,
            timeline=context,
        )
        system_entry = execute_codex_interaction(
            agent=system_cfg.get("agent", "coordinator"),
            role="system",
            client_name=system_cfg.get("client", "codex-cli"),
            capabilities=list(system_cfg.get("capabilities", [])),
            prompt=system_prompt,
            base_url=base_url,
            token=token,
            codex_bin=codex_bin,
            state=system_cfg.get("state", "online"),
            heartbeat_interval=heartbeat_interval,
            timeline_path=timeline_path,
            entry_type="system",
            extra_args=args.extra,
        )
        timeline.append(system_entry)
        if system_entry["returncode"] != 0:
            raise CLIError("System agent codex command failed; aborting schedule.")

        decision = extract_decision(system_entry["output"])
        agent_key = decision.get("agent")
        instruction = decision.get("instruction") or decision.get("prompt")
        if not agent_key or agent_key not in agents_cfg:
            raise CLIError(f"System agent selected unknown agent '{agent_key}'.")
        if not instruction:
            raise CLIError("System agent response missing 'instruction' field.")

        agent_cfg = agents_cfg[agent_key]
        agent_prompt_template = agent_cfg.get("prompt") or agent_cfg.get("prompt_template")
        if not agent_prompt_template:
            agent_prompt_template = "{instruction}"
        agent_prompt = agent_prompt_template.format(
            instruction=instruction,
            context=context,
            step=step,
            system_output=system_entry["output"],
            environment=environment_context,
            environment_context=environment_context,
        )
        agent_entry = execute_codex_interaction(
            agent=agent_key,
            role=agent_cfg.get("role", agent_key),
            client_name=agent_cfg.get("client", "codex-cli"),
            capabilities=list(agent_cfg.get("capabilities", [])),
            prompt=agent_prompt,
            base_url=base_url,
            token=token,
            codex_bin=codex_bin,
            state=agent_cfg.get("state", "online"),
            heartbeat_interval=heartbeat_interval,
            timeline_path=timeline_path,
            entry_type="agent",
            extra_args=args.extra,
        )
        timeline.append(agent_entry)
        if agent_entry["returncode"] != 0:
            raise CLIError(f"Agent '{agent_key}' codex command failed; aborting schedule.")

    print(f"Timeline written to {timeline_path}")


# ---------------------------------------------------------------------------
# Parser
# ---------------------------------------------------------------------------


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description="Agent-Society CLI")
    parser.set_defaults(func=None)

    sub = parser.add_subparsers(dest="command")

    login = sub.add_parser("login", help="Start a session and mark the agent online")
    login.add_argument("agent", help="Agent identifier (matches config/agents.json id)")
    login.add_argument("--client", default="codex-cli", help="Client label shown in the UI")
    login.add_argument("--capabilities", nargs="*", help="Optional capability tags")
    login.add_argument("--base-url", default=DEFAULT_BASE_URL, help="API base URL")
    login.add_argument("--token", help="Bearer token (defaults to AGENT_SOCIETY_TOKEN env)")
    login.set_defaults(func=cmd_login)

    heartbeat = sub.add_parser("heartbeat", help="Send a single heartbeat for the active session")
    heartbeat.add_argument("--state", default="online", choices=["online", "away"], help="Presence state to report")
    heartbeat.add_argument("--capabilities", nargs="*", help="Override capability set for this heartbeat")
    heartbeat.set_defaults(func=cmd_heartbeat)

    watch = sub.add_parser("watch", help="Send periodic heartbeats until stopped")
    watch.add_argument("--state", default="online", choices=["online", "away"], help="Presence state to report")
    watch.add_argument("--capabilities", nargs="*", help="Override capability set for heartbeats")
    watch.add_argument("--interval", type=int, default=60, help="Heartbeat interval in seconds")
    watch.set_defaults(func=cmd_watch)

    logout = sub.add_parser("logout", help="Close the current session and mark offline")
    logout.set_defaults(func=cmd_logout)

    presence = sub.add_parser("presence", help="List all agent presence records")
    presence.add_argument("--base-url", default=DEFAULT_BASE_URL, help="API base URL")
    presence.add_argument("--token", help="Bearer token (defaults to AGENT_SOCIETY_TOKEN env)")
    presence.add_argument("--format", choices=["table", "json"], default="table", help="Output format")
    presence.set_defaults(func=cmd_presence)

    timeline = sub.add_parser("timeline", help="Inspect orchestration timeline entries")
    timeline.add_argument("--base-url", default=DEFAULT_BASE_URL, help="API base URL")
    timeline.add_argument("--token", help="Bearer token (defaults to env)")
    timeline.add_argument("--limit", type=int, default=20, help="Number of entries to fetch")
    timeline.add_argument("--format", choices=["table", "json"], default="table", help="Output format")
    timeline.add_argument("--follow", action="store_true", help="Stream updates continuously")
    timeline.set_defaults(func=cmd_timeline)

    memories = sub.add_parser("memories", help="Show top-scored memories for an agent")
    memories.add_argument("--agent", help="Agent identifier (defaults to active session)")
    memories.add_argument("--base-url", help="API base URL (defaults to session or env config)")
    memories.add_argument("--token", help="Bearer token (defaults to session/env)")
    memories.add_argument("--limit", type=int, default=10, help="Number of memories to display")
    memories.add_argument("--query", help="Optional search query to bias relevance")
    memories.add_argument("--format", choices=["table", "json"], default="table", help="Output format")
    memories.set_defaults(func=cmd_memories)

    reflect = sub.add_parser("reflect", help="Generate and store a persona reflection via codex")
    reflect.add_argument("agent", help="Agent identifier to summarise")
    reflect.add_argument("--base-url", help="API base URL (defaults to env/config)")
    reflect.add_argument("--token", help="Bearer token (defaults to env/session)")
    reflect.add_argument("--codex-bin", default="codex", help="Path to codex executable")
    reflect.add_argument("--client", default="persona-reflector", help="Client label for presence")
    reflect.add_argument("--capabilities", nargs="*", help="Capabilities advertised during reflection")
    reflect.add_argument("--role", help="Optional role label (defaults to configured role)")
    reflect.add_argument("--state", default="away", choices=["online", "away"], help="Presence state")
    reflect.add_argument("--heartbeat-interval", type=int, default=90, help="Seconds between heartbeats")
    reflect.add_argument("--timeline", help="Append details to this JSONL timeline file")
    reflect.add_argument("--memories", type=int, default=12, help="Number of memories to include")
    reflect.add_argument("--query", help="Optional query bias for memory retrieval")
    reflect.add_argument("--format", choices=["json", "text"], default="json", help="Output format")
    reflect.add_argument("--no-cache", action="store_true", help="Skip writing persona_summaries.json")
    reflect.add_argument("--no-event", action="store_true", help="Skip appending a reflection event")
    reflect.add_argument("--extra", nargs=argparse.REMAINDER, help="Additional args passed to codex")
    reflect.set_defaults(func=cmd_reflect)

    environment = sub.add_parser("environment", help="Show the environment/task tree and natural-language context")
    environment.add_argument("--base-url", default=DEFAULT_BASE_URL, help="API base URL")
    environment.add_argument("--token", help="Bearer token (defaults to env)")
    environment.add_argument("--limit", type=int, default=5, help="Depth limit for tasks per agent/status")
    environment.add_argument("--format", choices=["text", "json"], default="text", help="Output format")
    environment.set_defaults(func=cmd_environment)

    codex_run = sub.add_parser("codex-run", help="Execute a single codex prompt with automatic presence management")
    codex_run.add_argument("agent", help="Agent identifier for presence tracking")
    codex_run.add_argument("prompt", help="Prompt passed to `codex -p`")
    codex_run.add_argument("--role", help="Optional human-readable role label")
    codex_run.add_argument("--client", default="codex-cli", help="Client label shown in presence pane")
    codex_run.add_argument("--capabilities", nargs="*", help="Capability tags advertised during the session")
    codex_run.add_argument("--base-url", default=DEFAULT_BASE_URL, help="API base URL")
    codex_run.add_argument("--token", help="Bearer token (defaults to AGENT_SOCIETY_TOKEN env)")
    codex_run.add_argument("--codex-bin", default="codex", help="Path to the codex executable")
    codex_run.add_argument(
        "--heartbeat-interval",
        type=int,
        default=90,
        help="Seconds between heartbeats while codex runs (0 disables)",
    )
    codex_run.add_argument("--state", default="online", choices=["online", "away"], help="Presence state")
    codex_run.add_argument("--timeline", help="Append execution details to this JSONL timeline file")
    codex_run.add_argument("--extra", nargs=argparse.REMAINDER, help="Additional args passed to codex")
    codex_run.set_defaults(func=cmd_codex_run)

    orchestrate = sub.add_parser(
        "codex-orchestrate",
        help="Coordinate multiple codex roles under a system agent and capture a timeline",
    )
    orchestrate.add_argument("--config", default="config/codex_schedule.json", help="Schedule config (JSON or YAML)")
    orchestrate.add_argument("--steps", type=int, help="Override number of iterations")
    orchestrate.add_argument("--base-url", help="API base URL (overrides config)")
    orchestrate.add_argument("--token", help="Bearer token (defaults to env/config)")
    orchestrate.add_argument("--timeline", help="Timeline output path (defaults to config or data/codex_timeline.jsonl)")
    orchestrate.add_argument("--codex-bin", help="Path to codex binary (overrides config)")
    orchestrate.add_argument("--heartbeat-interval", type=int, help="Seconds between heartbeats")
    orchestrate.add_argument("--context", type=int, help="Number of prior timeline entries shared with agents")
    orchestrate.add_argument("--extra", nargs=argparse.REMAINDER, help="Additional args passed to codex")
    orchestrate.set_defaults(func=cmd_codex_orchestrate)

    return parser


def main(argv: list[str] | None = None) -> int:
    parser = build_parser()
    args = parser.parse_args(argv)
    if not args.command:
        parser.print_help()
        return 1
    try:
        args.func(args)
        return 0
    except CLIError as exc:  # pragma: no cover - user-facing error path
        print(f"error: {exc}", file=sys.stderr)
        return 2
    except httpx.HTTPError as exc:  # pragma: no cover - network failure
        print(f"network error: {exc}", file=sys.stderr)
        return 3


if __name__ == "__main__":  # pragma: no cover
    raise SystemExit(main())
